/*
 * ParseFinder.java
 *
 * Created on July 10, 2006, 3:10 PM
 *
 */

package galronnlp.pcfg.unpacker;

import java.util.Iterator;

import galronnlp.pcfg.parser.Chart;
import galronnlp.pcfg.parser.ChartEntry;
import galronnlp.pcfg.parser.Edge;
import galronnlp.pcfg.parser.EntryTriple;

import galronnlp.util.Symbol;
import galronnlp.util.Tree;
import galronnlp.util.BinaryTree;
import galronnlp.util.Tuple;
import galronnlp.util.UnaryTree;
import galronnlp.util.LexicalTree;
import galronnlp.util.SortedLinkedList;

/**
 * This class contains the methods necessary to unpack the CKY chart. 
 *
 * @author Daniel A. Galron
 */
public class ParseFinder {
    
    Chart chart;
    Symbol startSymbol;
    int begin;
    int end;
    
    /** Creates a new instance of ParseFinder */
    public ParseFinder(Chart c, int begin, int end, Symbol startSymbol) {
        this.chart = c;
        this.startSymbol = startSymbol;
        this.begin = begin;
        this.end = end;
    }
    
    /**
     * Return the highest scoring parse (i.e. the parse with the highest probability
     * in the chart.)
     */
    public Tree getBestParse() {
        return this.findBestParse(this.begin, this.end, this.startSymbol);
    }
    
    /**
     * Return the highest scoring parses within some beam distance of the best parse
     */
    public SortedLinkedList<Tuple> getBestParses(double beam) {
        return this.findBestParses(this.begin, this.end, this.startSymbol, beam);
    }
    
    
    private Tree findBestParse(int b, int e, Symbol c) {
        // Get the entry spanning (b, e) with symbol c
        ChartEntry entry = chart.get(b, e, c);
        // Get the highest scoring entry in the chart (b, e)
        EntryTriple topEntry = entry.items().getLast();
        double p = ((Double)topEntry.get(0)).doubleValue();
        Edge edge = (Edge)topEntry.get(2);
        System.out.println(topEntry);
        // If the edge spanning (b, e) is a unary edge, then find the best parse
        // of the subtree, using the Symbol in the RHS as the root of the subtree.
        // Then, create a new Unary Tree with the RHS subtree as its child, and 
        // return the new subtree.
        if(edge.type() == edge.UNARY) {
            Tree RHS = this.findBestParse(b, e, (Symbol)edge.entry());
            Tree ret = new UnaryTree(c, RHS, p);
            return ret;
        // If the edge (b, e) is a binary edge, then find the best parse tree whose
        // root is the left Symbol of the RHS and the best parse tree whose root is
        // the right Symbol of the RHS, and construct and return the new tree whose
        // root is the category spanning (b, e)
        } else if(edge.type() == edge.BINARY) {
            int m = ((Integer)((Tuple)edge.entry()).get(0)).intValue();
            Symbol lRHS = (Symbol)((Tuple)edge.entry()).get(1);
            Symbol rRHS = (Symbol)((Tuple)edge.entry()).get(2);
            Tree l_RHS = this.findBestParse(b, m, lRHS);
            Tree r_RHS = this.findBestParse(m, e, rRHS);
            Tree ret = new BinaryTree(c, l_RHS, r_RHS, p);
            return ret;
        // Base case: if the edge spanning (b, e) is a lexical edge, just return
        // a new lexical tree containing the lexical entry
        } else if(edge.type() == edge.LEXICAL) {
            return new LexicalTree(c, (Symbol)edge.entry(), p);
        }
        else {
            System.err.println("here " + edge.type());
            System.err.println("ERROR! UNKNOWN EDGE TYPE");
            return null;
        }
        
    }
    private SortedLinkedList<Tuple> findBestParses(int b, int e, Symbol c, double beam) {
        // Get the entry spanning (b, e) with symbol c
        System.err.println("(" + b + ", " + e + ": " + c + ")");
        ChartEntry entry = chart.get(b, e, c);
        // Get the highest scoring entry in the chart (b, e)
        double highestProb = ((Double)entry.items().getHighest().get(0)).doubleValue();
        SortedLinkedList<Tuple> ret = new SortedLinkedList<Tuple>();
        // Iterate through all the chart entries
        for(Iterator<EntryTriple> i = entry.items().iterator(); i.hasNext(); ) {
            // get the candidate production spanning (b, e) 
            EntryTriple candidate = i.next();
            double p = ((Double)candidate.get(0)).doubleValue();
            // check to see if the probability is within the beam of highest scoring
            // production. If it is, then:
            if(p >= highestProb * beam) {
                // get the edge
                Edge edge = (Edge)candidate.get(2);
                // if the edge is a unary edge, then:
                if(edge.type() == edge.UNARY) {
                    // Find the all the parses within the beam of the highest scoring
                    // production of the child of this edge.
                    SortedLinkedList<Tuple> RHS = this.findBestParses(b, e, (Symbol)edge.entry(), beam);
                    // Iterate through all the subtrees found, and make a new unary tree
                    // for each subtree found, and add it to the return SortedLinkedList.
                    for(Iterator<Tuple> it = RHS.iterator(); it.hasNext(); ) {
                        Tuple next = it.next();
                        Tuple newTup = new Tuple(2);
                        Tree newTree = new UnaryTree(c, (Tree)next.get(1), p);
                        newTup.put(0, new Double(p));
                        newTup.put(1, newTree);
                        ret.insort(newTup);
                    }
                // Otherwise, if the edge is a binary edge, then:
                } else if(edge.type() == edge.BINARY) {
                    // get the point which divides the RHS from the edge
                    int m = ((Integer)((Tuple)edge.entry()).get(0)).intValue();
                    Symbol lRHS = (Symbol)((Tuple)edge.entry()).get(1);
                    Symbol rRHS = (Symbol)((Tuple)edge.entry()).get(2);
                    // Find the best parses of the subtree whose root is the left
                    // Symbol of the RHS and of the subtree whose root is the right
                    // Symbol of the RHS
                    SortedLinkedList<Tuple> l_RHS = this.findBestParses(b, m, lRHS, beam);
                    SortedLinkedList<Tuple> r_RHS = this.findBestParses(m, e, rRHS, beam);
                    // Construct a subtree for each combination of l_RHS and r_RHS;
                    // that is, a tree for each l_RHS X r_RHS
                    for(Iterator<Tuple> il = l_RHS.iterator(); il.hasNext(); ) {
                        Tuple l_hand = il.next();
                        for(Iterator<Tuple> ir = r_RHS.iterator(); ir.hasNext(); ) {
                            Tuple r_hand = ir.next();
                            Tuple newTup = new Tuple(2);
                            Tree newTree = new BinaryTree(c, (Tree)l_hand.get(1), (Tree)r_hand.get(1), p);
                            newTup.put(0, new Double(p));
                            newTup.put(1, newTree);
                            ret.insort(newTup);
                        }
                    }
                // Base case: return all lexical trees within the beam of the highest
                // scoring lexical entry.
                } else if(edge.type() == edge.LEXICAL) {
                    Tuple newTup = new Tuple(2);
                    Tree newTree = new LexicalTree(c, (Symbol)edge.entry(), p);
                    newTup.put(0, new Double(p));
                    newTup.put(1, newTree);
                    ret.insort(newTup);
                } else {
                    System.err.println("HERE! " + edge);
                    System.err.println("ERROR! UNKNOWN EDGE TYPE");
                    return null;
                }                
            } 
        }
        return ret;
    }
}
